#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
安装NLP依赖
检查并安装必要的自然语言处理包
"""

import sys
import subprocess
import importlib
import importlib.util
import logging
import os
import platform
from distutils.version import LooseVersion
from datetime import datetime
from pathlib import Path
import time

# 设置日志
log_dir = "logs"
if not os.path.exists(log_dir):
    os.makedirs(log_dir)
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
log_file = os.path.join(log_dir, f"nlp_dependencies_install_{timestamp}.log")
logging.basicConfig(level=logging.INFO,
                   format='%(asctime)s - %(levelname)s - %(message)s',
                   handlers=[
                       logging.FileHandler(log_file, encoding='utf-8'),
                       logging.StreamHandler()
                   ])

logger = logging.getLogger(__name__)

# 检查Python版本
def check_python_version():
    """检查Python版本，返回版本号和兼容性状态"""
    version_info = sys.version_info
    version_str = f"{version_info.major}.{version_info.minor}.{version_info.micro}"
    logging.info(f"检测到Python版本: {version_str}")
    
    if version_info.major < 3 or (version_info.major == 3 and version_info.minor < 7):
        logging.error(f"Python版本 {version_str} 过低。需要Python 3.7或更高版本。")
        return version_str, False
    
    # Python 3.12兼容性检查
    if version_info.major == 3 and version_info.minor >= 12:
        logging.warning(f"Python {version_str} 可能与某些NLP库存在兼容性问题")
        logging.warning("将尝试安装兼容版本，但可能无法保证所有功能正常")
        return version_str, True
    
    logging.info(f"Python版本 {version_str} 兼容")
    return version_str, True

# 检查包是否已安装
def is_package_installed(package_name):
    """检查包是否已安装"""
    return importlib.util.find_spec(package_name) is not None

# 获取包版本
def get_package_version(package_name):
    """获取包的安装版本"""
    try:
        result = subprocess.run(
            [sys.executable, "-m", "pip", "show", package_name], 
            capture_output=True, 
            text=True
        )
        if result.returncode == 0:
            for line in result.stdout.split('\n'):
                if line.startswith('Version:'):
                    return line.split(':', 1)[1].strip()
        return None
    except Exception as e:
        logging.error(f"获取{package_name}版本时出错: {str(e)}")
        return None

# 安装包
def install_package(package_name, is_optional=False, py312_alternative=None):
    """
    安装指定包，考虑Python 3.12兼容性
    
    Args:
        package_name: 要安装的包名
        is_optional: 是否为可选包（安装失败不会报错）
        py312_alternative: Python 3.12的替代包
    """
    try:
        # 检查包是否已安装
        if is_package_installed(package_name.split('==')[0].split('[')[0]):
            logging.info(f"已安装: {package_name}")
            return True
        
        # 针对Python 3.12的处理
        py_version = sys.version_info
        if py_version.major == 3 and py_version.minor >= 12 and py312_alternative:
            logging.info(f"检测到Python 3.12+，将使用替代包: {py312_alternative}")
            package_name = py312_alternative
        
        # 安装包
        logging.info(f"正在安装: {package_name}")
        cmd = [sys.executable, "-m", "pip", "install", "--no-cache-dir", package_name]
        result = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
        
        if result.returncode != 0:
            if is_optional:
                logging.warning(f"可选包 {package_name} 安装失败: {result.stderr}")
                return False
            else:
                logging.error(f"必需包 {package_name} 安装失败: {result.stderr}")
                return False
        
        logging.info(f"成功安装: {package_name}")
        return True
    except Exception as e:
        if is_optional:
            logging.warning(f"可选包 {package_name} 安装过程出错: {str(e)}")
            return False
        else:
            logging.error(f"必需包 {package_name} 安装过程出错: {str(e)}")
            return False

# 安装fuzzy匹配功能
def install_fuzzy_matching():
    """安装模糊匹配依赖"""
    logger.info("检查模糊匹配依赖...")
    
    if get_package_version("fuzzywuzzy"):
        logger.info("fuzzywuzzy 已安装")
        fuzzy_status = True
    else:
        fuzzy_status = install_package("fuzzywuzzy[speedup]")
    
    if fuzzy_status and not get_package_version("python-Levenshtein"):
        logger.info("正在安装 python-Levenshtein 以提高模糊匹配性能...")
        levenshtein_status = install_package("python-Levenshtein")
        if not levenshtein_status:
            logger.warning("python-Levenshtein 安装失败，将使用较慢的纯Python实现")
    
    return fuzzy_status

# 安装中文分词
def install_chinese_tokenization():
    """安装中文分词依赖"""
    logger.info("检查中文分词依赖...")
    
    if get_package_version("jieba"):
        logger.info("jieba 已安装")
        return True
    else:
        logger.info("正在安装 jieba 中文分词库...")
        return install_package("jieba")

# 安装spaCy和中文模型
def install_spacy():
    """安装spaCy及其语言模型"""
    logger.info("检查spaCy依赖...")
    
    py_version = platform.python_version()
    major, minor, _ = map(int, py_version.split('.'))
    
    # Python 3.12 兼容性处理
    if major == 3 and minor >= 12:
        logger.warning("检测到Python 3.12+，将使用兼容模式安装spaCy")
        spacy_version = "3.7.2"  # 或其他已知与Python 3.12兼容的版本
        spacy_package = f"spacy=={spacy_version}"
    else:
        spacy_package = "spacy"
    
    if get_package_version("spacy"):
        logger.info(f"spaCy 已安装: {get_package_version('spacy')}")
        spacy_status = True
    else:
        logger.info(f"正在安装 {spacy_package}...")
        spacy_status = install_package(spacy_package, disable_pip_version_check=True)
    
    if not spacy_status:
        logger.error("spaCy 安装失败，命名实体识别功能将不可用")
        return False
    
    # 安装英文模型
    en_model = "en_core_web_sm"
    try:
        logger.info(f"检查 {en_model} 模型...")
        subprocess.run([sys.executable, "-m", "spacy", "validate"], capture_output=True, check=True)
        
        # 检查模型是否已安装
        result = subprocess.run([sys.executable, "-m", "spacy", "info"], capture_output=True, text=True)
        if en_model not in result.stdout:
            logger.info(f"正在下载 {en_model} 模型...")
            subprocess.run([sys.executable, "-m", "spacy", "download", en_model], check=True)
            logger.info(f"{en_model} 模型安装完成")
        else:
            logger.info(f"{en_model} 模型已安装")
        
        return True
    except subprocess.CalledProcessError:
        logger.error(f"安装 {en_model} 模型失败")
        return False

# 安装正则表达式增强库
def install_regex():
    if is_package_installed("regex"):
        version = get_package_version("regex")
        logger.info(f"regex 已安装 (版本: {version})")
        return True
    else:
        if install_package("regex"):
            logger.info("成功安装regex增强库")
            return True
        else:
            logger.warning("regex安装失败")
            return False

# 安装终端颜色库
def install_colorama():
    if is_package_installed("colorama"):
        version = get_package_version("colorama")
        logger.info(f"colorama 已安装 (版本: {version})")
        return True
    else:
        if install_package("colorama"):
            logger.info("成功安装colorama终端颜色库")
            return True
        else:
            logger.warning("colorama安装失败")
            return False

def install_additional_libraries():
    """安装其他辅助库"""
    libraries = {
        "regex": install_regex(),
        "colorama": install_colorama(),
        "nltk": install_package("nltk")
    }
    
    return all(libraries.values())

def check_installations():
    """检查所有依赖的安装状态"""
    packages = {
        "fuzzywuzzy": get_package_version("fuzzywuzzy"),
        "python-Levenshtein": get_package_version("python-Levenshtein"),
        "jieba": get_package_version("jieba"),
        "spacy": get_package_version("spacy"),
        "regex": get_package_version("regex"),
        "colorama": get_package_version("colorama"),
        "nltk": get_package_version("nltk")
    }
    
    logger.info("安装状态摘要:")
    for package, version in packages.items():
        status = "已安装" if version else "未安装"
        version_info = f" (版本: {version})" if version else ""
        logger.info(f"- {package}: {status}{version_info}")
    
    return all(version is not None for version in packages.values())

def create_startup_script():
    """创建启动脚本"""
    logging.info("创建启动脚本...")
    
    if platform.system() == "Windows":
        # Windows批处理脚本
        script_path = "start_data_clean.bat"
        with open(script_path, "w") as f:
            f.write('@echo off\n')
            f.write('echo 启动数据清洗工具...\n')
            f.write('python main.py --mode clean\n')
            f.write('pause\n')
    else:
        # Unix/Linux/Mac shell脚本
        script_path = "start_data_clean.sh"
        with open(script_path, "w") as f:
            f.write('#!/bin/bash\n')
            f.write('echo "启动数据清洗工具..."\n')
            f.write('python main.py --mode clean\n')
            f.write('read -p "按回车键继续..."\n')
        
        # 设置执行权限
        try:
            os.chmod(script_path, 0o755)
        except Exception as e:
            logging.warning(f"无法设置脚本执行权限: {str(e)}")
    
    logging.info(f"启动脚本已创建: {script_path}")
    return script_path

def main():
    """安装NLP依赖并创建启动脚本"""
    logging.info("===== NLP依赖安装程序开始 =====")
    
    # 检查Python版本
    version_str, is_compatible = check_python_version()
    if not is_compatible:
        logging.error("Python版本不兼容，停止安装")
        return False
    
    # 基本依赖
    required_packages = [
        "pandas",
        "numpy",
        "openpyxl",
        "python-dateutil",
        "tqdm"
    ]
    
    # NLP专用依赖
    nlp_packages = [
        "nltk",
        "fuzzywuzzy",
        "python-Levenshtein",  # 用于fuzzywuzzy加速
    ]
    
    # 中文NLP支持
    chinese_nlp_packages = [
        "jieba", 
    ]
    
    # 可选依赖
    optional_packages = [
        ("spacy", True, "spacy<3.7.0"),  # Python 3.12下使用较老版本
        ("zh_core_web_sm", True, None),  # spaCy模型
        ("transformers", True, None),    # 可能有用但非必需
    ]
    
    # 安装基本依赖
    logging.info("安装基本依赖...")
    basic_success = all(install_package(pkg) for pkg in required_packages)
    
    # 安装NLP依赖
    logging.info("安装NLP依赖...")
    nlp_success = all(install_package(pkg) for pkg in nlp_packages)
    
    # 安装中文NLP依赖
    logging.info("安装中文NLP依赖...")
    chinese_success = all(install_package(pkg) for pkg in chinese_nlp_packages)
    
    # 安装可选依赖
    logging.info("安装可选NLP依赖...")
    for pkg, is_optional, py312_alt in optional_packages:
        install_package(pkg, is_optional, py312_alt)
    
    # 创建启动脚本
    script_path = create_startup_script()
    
    # 总结
    if basic_success and nlp_success and chinese_success:
        logging.info("===== 所有必需依赖安装成功 =====")
        return True
    else:
        logging.warning("===== 部分依赖安装失败，查看日志获取详情 =====")
        return False

if __name__ == "__main__":
    try:
        success = main()
        if not success:
            sys.exit(1)
    except Exception as e:
        logging.error(f"安装过程发生错误: {str(e)}")
        sys.exit(1) 